from bert4keras.models import build_transformer_model
from bert4keras.backend import keras, set_gelu
from keras.layers import Lambda, Dense, Dropout
from keras.regularizers import l2

set_gelu('tanh')  # 切换gelu版本


def build_model(config_path, checkpoint_path, class_nums):
    bert = build_transformer_model(
        config_path=config_path,
        checkpoint_path=checkpoint_path,
        return_keras_model=False)

    # BERT模型在文本前插入一个[CLS]符号，并将该符号对应的输出向量作为整篇文本的语义表示，用于文本分类
    cls_features = Lambda(lambda x: x[:, 0])(bert.model.output)

    # 防止过拟合
    Dropout(0.5)(cls_features)

    dense = Dense(
        units=class_nums,
        activation='softmax',
        kernel_initializer=bert.initializer,
        kernel_regularizer=l2(0.0003)  # 防止过拟合
    )(cls_features)

    model = keras.models.Model(bert.model.input, dense)
    return model
